import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

from skimage.segmentation import clear_border
from skimage.measure import label, regionprops
from skimage.color import label2rgb

import numpy as np
from PIL import Image

def calcLabelImage(image):
    cleared = clear_border(image)

    label_image = label(cleared)
    
    return label_image

def drawAreaRect(predict_y, limit=10):
    threshold = predict_y.max() - 0.9 * (predict_y.max() - predict_y.min())
    image = (predict_y > threshold).astype('uint8')

    label_image = calcLabelImage(image)
    image_label_overlay = label2rgb(label_image, image=image)

    fig, ax = plt.subplots(figsize=(10, 6))
    ax.imshow(image_label_overlay)

    for region in regionprops(label_image):
        if region.area >= limit:
            minr, minc, maxr, maxc = region.bbox
            rect = mpatches.Rectangle((minc, minr), maxc - minc, maxr - minr,
                                      fill=False, edgecolor='red', linewidth=2)
            ax.add_patch(rect)

    plt.show()

def calcAreaBoxs(predict_y, limit=10):
    threshold = predict_y.max() - 0.9 * (predict_y.max() - predict_y.min())
    image = (predict_y > threshold).astype('uint8')

    label_image = calcLabelImage(image)

    boxs = []
    for region in regionprops(label_image):
        if region.area >= limit:
            minr, minc, maxr, maxc = region.bbox
            boxs.append([[minr, minc], [maxr, maxc]])
    return boxs

def formatImage(img, size=(128, 128)):
    image = img
    fg_w, fg_h = image.size
    bg_w, bg_h = size
    
    if image.mode == 'L':
        out_image = np.zeros((bg_h, bg_w), dtype='uint8') + 255
    elif image.mode == 'RGB':
        out_image = np.zeros((bg_h, bg_w, 3), dtype='uint8') + 255
    elif image.mode == 'RGBA':
        out_image = np.zeros((bg_h, bg_w, 4), dtype='uint8') + 255

    if fg_h == fg_w:
        image2 = image.resize(size, resample=Image.BOX)
        out_image = np.array(image2, dtype='uint8')
    elif fg_h > fg_w:
        h = bg_h
        w = int(h / fg_h * fg_w)
        image2 = image.resize((w, h), resample=Image.BOX)
        image2_arr = np.array(image2, dtype='uint8')
        out_image[:, (bg_w - w) // 2: (bg_w - w) // 2 + w] = image2_arr
    elif fg_w > fg_h:
        w = bg_w
        h = int(w / fg_w * fg_h)
        image2 = image.resize((w, h), resample=Image.BOX)
        image2_arr = np.array(image2, dtype='uint8')
        out_image[(bg_h - h) // 2: (bg_h - h) // 2 + h, :] = image2_arr
    return out_image

def calcSourceBoxs(image, boxs, a):
    if len(boxs) == 0: return np.array([])

    w, h = image.size
    size = np.array([h, w])
    k = np.max(size) / a

    boxs_arr = np.array(boxs)
    boxs_arr = boxs_arr * k + (size - np.max(size)) / 2
    return np.maximum(0, np.round(boxs_arr)).astype('int')

def calcSplitPTS(Y, flag=True):
    i = 0
    split_pts = []
    for y in Y:
        if (y > 0) == flag:
            split_pts.append(i)
            flag = not flag
        i += 1
    return split_pts

def calcAccurateBoxs(image_arr, boxs_arr):
    boxs_arr_int = np.maximum(0, boxs_arr).astype('int')

    for box_arr in boxs_arr_int:
        if np.min(box_arr[1, :] - box_arr[0, :]) < 1: continue
        sub_area = image_arr[box_arr[0, 0]:box_arr[1, 0], box_arr[0, 1]:box_arr[1, 1]]
        pix_min, pix_max = sub_area.min(), sub_area.max()
        normal_area = 1 - (sub_area.astype('float32') - pix_min) / (pix_max - pix_min + 1e-06)
        
        ext_size = 3
        
        Yc = normal_area.mean(axis=1)
        Yc = (Yc > Yc.mean()).astype('int')
        split_pts_c = calcSplitPTS(Yc)
        if len(split_pts_c) < 2: continue
        split_top = split_pts_c[0] - ext_size if split_pts_c[0] - ext_size > 0 else 0
        split_bottom = split_pts_c[-1] + ext_size
        normal_area_c = normal_area[split_top:split_bottom, :]
        
        Yr = normal_area_c.max(axis=0)
        Yr = (Yr > 0.25 * Yr.max()).astype('int')
        split_pts_r = calcSplitPTS(Yr, Yr[0] < 1)
        if len(split_pts_r) < 8: continue
        dists = np.convolve(split_pts_r, [1, -1], mode='valid')
        split_left = split_pts_r[0] - ext_size
        split_right = split_pts_r[-1] + ext_size if Yr[-1] < 1 else Yr.size + 1
        if dists[0::2][1:-1].mean() < dists[1::2][1:-1].mean(): # mode 1010
            if dists[0::2][0] > dists[0::2][1:].max(): # mode 1001
                split_left = split_pts_r[1] - ext_size
            else:
                split_left = 0
        else:
            if dists[1::2][0] > 1.5 * dists[1::2][1:].max(): # mode 0100
                split_left = split_pts_r[2] - ext_size
        
        box_arr[1, 0] = box_arr[0, 0] + split_bottom
        box_arr[1, 1] = box_arr[0, 1] + split_right
        box_arr[0, 0] += split_top
        box_arr[0, 1] += split_left

    return boxs_arr_int
